package com.han.gateway.filter;

import com.alibaba.fastjson2.JSONObject;
import com.han.gateway.util.XSSUtil;
import io.micrometer.common.util.StringUtils;
import lombok.AllArgsConstructor;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

@Component
@AllArgsConstructor
public class XssFilter implements GlobalFilter {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest serverHttpRequest = exchange.getRequest();
        HttpMethod method = serverHttpRequest.getMethod();

        if (method == HttpMethod.GET) {
            return dealForGet(exchange, chain);
        }

        if (method == HttpMethod.POST) {
            return dealForPost(exchange, chain, serverHttpRequest);
        }

        return chain.filter(exchange);
    }

    private static Mono<Void> dealForGet(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI uri = exchange.getRequest().getURI();
        String originalQuery = uri.getRawQuery();

        if (StringUtils.isBlank(originalQuery)) {
            return chain.filter(exchange);
        }

        // 把originalQuery中的xss脚本删除掉
        String replaceQuery = XSSUtil.stripXSS(originalQuery);
        try {
            URI newUri = UriComponentsBuilder.fromUri(uri)
                    .replaceQuery(replaceQuery)
                    .build(ServerWebExchangeUtils.containsEncodedParts(uri))
                    .toUri();

            // 重写uri
            ServerHttpRequest request = exchange.getRequest().mutate().uri(newUri).build();
            return chain.filter(exchange.mutate().request(request).build());
        } catch (RuntimeException ex) {
            throw new IllegalStateException("Invalid URI query: \"" + originalQuery + "\"", ex);
        }
    }

    private static Mono<Void> dealForPost(ServerWebExchange exchange, GatewayFilterChain chain, ServerHttpRequest serverHttpRequest) {
        ServerHttpRequest requestDecorator = new ServerHttpRequestDecorator(serverHttpRequest) {
            @Override
            public Flux<DataBuffer> getBody() {
                return serverHttpRequest.getBody()
                        .map(dataBuffer -> {
                            ByteBuffer byteBuffer = ByteBuffer.allocate(dataBuffer.capacity());
                            dataBuffer.toByteBuffer(byteBuffer);
                            String body = new String(byteBuffer.array(), StandardCharsets.UTF_8);

                            JSONObject jsonObject = JSONObject.parseObject(body);
                            JSONObject newJsonObject = new JSONObject();
                            jsonObject.forEach((k, v) -> {
                                newJsonObject.put(k, XSSUtil.stripXSS(v.toString()));
                            });

                            String encodedBody = newJsonObject.toString();
                            // 将编码后的内容重新封装成DataBuffer
                            return DefaultDataBufferFactory.sharedInstance.wrap(encodedBody.getBytes(StandardCharsets.UTF_8));
                        });
            }
        };
        return chain.filter(exchange.mutate().request(requestDecorator).build());
    }
}
